import math
import subprocess

import numpy as np
import torch

import whisper
from ai.utils import utils_showing, utils_file
from whisper import audio, timing

audio_file_path = "E:\gengxuelong_study\server_local_adapter\\ai\data\small_aishell\dev\BAC009S0724W0121.wav"


def do_test_load_audio():
    audio_file_path = "E:\gengxuelong_study\server_local_adapter\\ai\data\small_aishell\dev\BAC009S0724W0121.wav"
    # audio_file_path = "/home/work_nfs/common/data/data_aishell/wav/dev/S0724/BAC009S0724W0121.wav"
    wav_data = audio.load_audio(audio_file_path)
    print(wav_data)


def do_test_save_mel_filter():
    # audio.save_mel_filters()
    dic = np.load("./output/gxl.npz")
    print(type(dic))
    print(dic['name'])
    print(dic['mel_80'].shape)


def do_test_get_mel_spec():
    # audio_file_path = "/home/work_nfs/common/data/data_aishell/wav/dev/S0724/BAC009S0724W0121.wav"
    wav_data = audio.log_mel_spectrogram(audio_file_path)
    print(wav_data.shape)


def do_test_median():
    x = torch.tensor([[1, 2, 3, 4, 5, 6, 7], [4, 5, 6, 7, 8, 9, 10], [7, 8, 9, 10, 11, 12, 13]])
    median_values = torch.median(x, dim=1)
    filter_width = 3
    # sort()函数返回一个特元组, 第一个值为value, 第二个值为index
    median_values2 = x.unfold(-1, filter_width, 1).sort()[0][..., filter_width // 2]
    print(median_values2)


def do_test_tiktoken():
    """"""
    import tiktoken
    encode = tiktoken.get_encoding("gpt4")
    codes = encode.encode("Hello, world! 我是耿雪龙")
    chars = encode.decode(codes)
    char_list = [encode.decode([i]) for i in codes]
    print(chars)
    print(char_list)


def do_test_timing():
    """"""
    range_x = np.arange(0, 10, 0.01)
    print(range_x)
    print(len(range_x))
    sin_y = np.sin(range_x)
    print(sin_y)
    # utils_showing.show_lines(range_x,sin_y)
    sin_y2 = whisper.timing.median_filter(torch.from_numpy(sin_y), 11)
    utils_showing.show_lines(range_x, sin_y2)


import torch
from torch.distributions import Categorical


def do_test_category():
    # 假设我们有一个策略网络，它输出两个动作的概率
    probs = torch.tensor([0.4, 0.6])
    # 创建一个 Categorical 分布对象
    m = Categorical(probs)
    # 从分布中采样一个动作
    action = m.sample()
    # 打印采样结果
    print(action)
    # 输出：tensor(1)
    # 计算采样动作的对数概率
    log_prob = m.log_prob(action)
    print(np.log2(0.4))
    print(np.log2(0.6))
    print(np.power(0.6, -0.5108))
    # 打印对数概率
    print(log_prob)
    # 输出：tensor(-0.5108)


def do_test_drawing_func():
    def func(x):
        return x ** 0.4

    utils_showing.show_func(func)


def do_test_argmax_shape():
    x = torch.randn(10, 1, 5)
    y = x.argmax(dim=-1)
    print(y.shape)


def do_common():
    """"""
    tokens = torch.randn(3, 5)
    tokens = tokens.repeat_interleave(8, dim=0)
    print(tokens)


import whisper


def do_test_whisper_gxl():
    print(whisper.available_models())
    whisper_model = whisper.load_model("tiny", "cpu", "./output")
    print(whisper_model)
    print(whisper_model.dims)
    token_fake_input = torch.randint(0, 51865, (1, 448))
    audio_fake_input = torch.randn(1, 80, 3000)
    print(whisper_model(audio_fake_input, token_fake_input).shape)


def do_test_use():
    """"""
    cwd = "E:\gengxuelong_study\server_local_adapter\\ai"
    subprocess.run(['python', '-m', 'whisper'], cwd=cwd)


def do_download():
    """"""
    gxl_loader = utils_file.GxlDownloader("./output/gxl_download/")
    gxl_loader.set_suffix('png')
    gxl_loader.download(
        "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcT7ctJE5wstB8huhbOdrFPQaZEPfZRTGpMi0g&usqp=CAU")



if __name__ == '__main__':
    """"""
    # do_common()
    # do_test_whisper_gxl()
    do_test_use()
    # do_download()
